import torch
from itertools import chain

# Load the image features and IDs
image_features_path = "/data/jzw/MSCOCO/features/coco_val2017_image_features.pt"
image_ids_path = "/data/jzw/MSCOCO/features/coco_val2017_image_ids.pt"  # Assuming IDs are stored in a separate file
output_features_path = "deduplicated_image_features.pkl"
output_ids_path = "deduplicated_image_ids.pkl"

# Load the features and IDs
image_features = torch.load(image_features_path)
image_ids = torch.load(image_ids_path)

# # Expand the loop for debugging image_ids
# for i, batch in enumerate(image_ids):
#     print(f"Processing batch {i}: {batch}")
#     if isinstance(batch, int):
#         print(f"Batch at image_ids[{i}] is an int: {batch}")
#         batch = [batch]  # Convert int to list
#     else:
#         print(f"Batch at image_ids[{i}] is a list: {batch}")
#         batch = list(chain.from_iterable(batch))
#     print(f"Processed batch {i}: {batch}")

# Concatenate all batches in image_ids into a single list
image_ids = list(chain.from_iterable(image_ids))

assert len(image_ids) == len(image_features), "Mismatch between number of features and IDs"

# Create a dictionary to deduplicate features based on IDs
unique_features = {}
for feature, image_id in zip(image_features, image_ids):
    unique_features[image_id] = feature  # Overwrite duplicates with the latest feature

# Extract deduplicated features and IDs
deduplicated_ids = list(unique_features.keys())
deduplicated_features = torch.stack(list(unique_features.values()))

# Save the deduplicated features and IDs
torch.save(deduplicated_features, output_features_path)
torch.save(deduplicated_ids, output_ids_path)
print(f"Deduplicated features saved to {output_features_path}")
print(f"Deduplicated IDs saved to {output_ids_path}")